/*
  Copyright (C) 2011, Kyungjoo Kim
  
  This file is part of UHM (Unassembled HyperMatrix) sparse direct solver.
  
  All rights reserved.

  This library is free software; you can redistribute it and/or
  modify it under the terms of the GNU Lesser General Public
  License as published by the Free Software Foundation; either
  version 2.1 of the License, or (at your option) any later version.
  
  This library is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  Lesser General Public License for more details.
  
  You should have received a copy of the GNU Lesser General Public
  License along with this library; if not, write to the Free Software
  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
  Also add information on how to contact you by electronic and paper mail.
  
  
  Kyungjoo Kim
  iamkyungjoo@gmail.com
*/
#ifndef UHM_PLUGIN_PETSC_HXX
#define UHM_PLUGIN_PETSC_HXX

namespace uhm {

  typedef class Net_*     Net;
  typedef class Mesh_*    Mesh;
  typedef class Node_*    Node;
  typedef class Element_* Element;
  typedef class Matrix_*  Matrix;

  namespace plugin {

    typedef class Petsc_* Petsc;

    // Petsc return for success is zero.
    class Petsc_ : public Net_ {
      
    protected:
      PetscMPIInt nrank, irank;

      int datatype, n_dof, n_rhs, n_nz;

      Mat A, B, X, R;
      Vec    b, x, r;

      KSP ksp;

      double tol;

      // control flags
      bool spd, verbose_col_op, verbose;

      void _init(int datatype, int spd, int verbose, int verbose_col_op);

    public:
      Petsc_();
      Petsc_(int datatype);
      Petsc_(int datatype, int spd);
      virtual ~Petsc_();
      
      int obj_view() { return this->obj_view(PETSC_COMM_WORLD, stdout); }
      int obj_view(MPI_Comm comm, FILE *stream);

      bool operator<(const Petsc_ &b) const;
      
      virtual void lock();
      virtual void unlock();
      
      bool is_complex();
      bool is_spd();
      
      void set_ksp(KSP ksp);
      void set_tolerence(double tol);
      
      int get_comm_rank();
      int get_comm_size();

      int get_petsc_datatype();
      
      void set_n_rhs(int n_rhs);

      int get_n_dof();
      int get_n_rhs();
      int get_n_nz();

      double get_tolerence();

      int initialize(int argc, char **argv);
      int finalize();

      int create_ksp();
      int create_ksp(MatStructure flag);
      int delete_ksp();
      int set_pc_on_ksp(PCType type);

      int solve_dist();
      int check_dist();

      int condition(PetscReal &cond);
      int view(int obj, PetscViewer viewer);

      int get_mat(int obj, Mat *M);
      int get_vec(int obj, Vec *v);

      int convergence();
      int convergence(MPI_Comm comm, FILE *stream);

      int create_matrix();
      int delete_matrix();
      
      int set_values_on_rhs_begin(Mat B, Vec b, int loc);
      int set_values_on_rhs_end  (Mat B, Vec b, int loc);

    };
    
    inline Petsc_::Petsc_()             { 
      this->_init(UHM_REAL, false, false, false); 
    }
    inline Petsc_::Petsc_(int datatype) { 
      this->_init(datatype, false, false, false); 
    }
    inline Petsc_::Petsc_(int datatype, int spd) { 
      this->_init(datatype, spd, false, false); 
    }
    inline Petsc_::~Petsc_() { }

    inline void   Petsc_::_init(int datatype, int spd, 
                                int verbose, int verbose_col_op) {
      switch (datatype) {
      case UHM_REAL:    this->datatype = UHM_REAL;    break;
      case UHM_COMPLEX: this->datatype = UHM_COMPLEX; break;
      default:
        UHM_ERROR(false,
                  ">> Datatype is not supported");
      }
      this->tol              = 1.0e-10;
      this->spd              = (bool)spd;
      this->verbose          = (bool)verbose;
      this->verbose_col_op   = (bool)verbose_col_op;
    }

#undef  __FUNCT__
#define __FUNCT__ "uhm::plugin::Petsc_::obj_view"
    inline int    Petsc_::obj_view(MPI_Comm comm, FILE *stream) {
      PetscErrorCode ierr;
      ierr = PetscPrintf(PETSC_COMM_WORLD,
                         "Petsc Plugin to UHM\n");CHKERRQ(ierr);
      return (0);
    }

    inline bool   Petsc_::is_complex() { return this->datatype == UHM_COMPLEX; }
    inline bool   Petsc_::is_spd() { return this->spd; }

    inline void   Petsc_::set_ksp(KSP ksp) { this->ksp = ksp; }
    inline void   Petsc_::set_tolerence(double tol) { this->tol = tol; }

    inline int    Petsc_::get_comm_rank() { return this->irank; }
    inline int    Petsc_::get_comm_size() { return this->nrank; }

    inline int    Petsc_::get_petsc_datatype() {
      if (this->is_complex()) return PETSC_COMPLEX;
      return PETSC_DOUBLE;
    }

    inline void   Petsc_::set_n_rhs(int n_rhs) { this->n_rhs = n_rhs; }

    inline int    Petsc_::get_n_dof()     { return this->n_dof; }
    inline int    Petsc_::get_n_rhs()     { return this->n_rhs; }
    inline int    Petsc_::get_n_nz()      { return this->n_nz; }
    inline double Petsc_::get_tolerence() { return this->tol; }

#undef  __FUNCT__
#define __FUNCT__ "uhm::plugin::Petsc_::initialize"
    inline int    Petsc_::initialize(int argc, char **argv) {
      PetscErrorCode ierr;
      PetscBool      
        spd = PETSC_FALSE, verbose = PETSC_TRUE, verbose_col_op = PETSC_TRUE;
      PetscInt 
        method = UHM_LU_PIV;

      char help[] = "UHM plugin to Petsc KSP solver\n";
      ierr = PetscInitialize(&argc, &argv, (char*)0, help);CHKERRQ(ierr); 
      
      ierr = MPI_Comm_rank(PETSC_COMM_WORLD, &this->irank);CHKERRQ(ierr);
      ierr = MPI_Comm_size(PETSC_COMM_WORLD, &this->nrank);CHKERRQ(ierr);
     
      ierr = PetscOptionsBegin(PETSC_COMM_WORLD, 
                               "", 
                               "uhm::plugin::petsc options", 
                               "KSP");CHKERRQ(ierr);
      ierr = PetscOptionsBool("-petsc_spd", "use 0) GMRESuse 1) CG", 
                              __FILE__, spd, &spd,
                              PETSC_NULL);CHKERRQ(ierr);
      ierr = PetscOptionsBool("-petsc_verbose", 
                              "verbose output 0) disable 1) enable",
                              __FILE__, verbose, &verbose,
                              PETSC_NULL);CHKERRQ(ierr);
      ierr = PetscOptionsBool("-petsc_verbose_col_op", 
                              "0) hide or 1) show operation on each rhs vector",
                              __FILE__, verbose_col_op, &verbose_col_op,
                              PETSC_NULL);CHKERRQ(ierr);
      ierr = PetscOptionsInt("-petsc_uhm_decompose", 
                             "decomposition method for UHM", 
                             __FILE__, method, &method, PETSC_NULL);CHKERRQ(ierr);
      ierr = PetscOptionsEnd();CHKERRQ(ierr);

      this->spd            = (bool)spd;
      this->verbose        = (bool)verbose;
      this->verbose_col_op = (bool)verbose_col_op;
      this->get_self()->set_decomposition_method(method);

      if (this->verbose) {
        ierr = PetscPrintf(PETSC_COMM_SELF,
                           ">> Irank %d, Nrank %d\n", 
                           this->irank, this->nrank);CHKERRQ(ierr);
      }

      return (0);
    }

#undef  __FUNCT__
#define __FUNCT__ "uhm::plugin::Petsc_::finalize"
    inline int    Petsc_::finalize() {
      PetscErrorCode ierr;
      ierr = PetscFinalize();

      return (0);
    }

#undef  __FUNCT__
#define __FUNCT__ "uhm::plugin::Petsc_::create_ksp"
    inline int    Petsc_::create_ksp() {
      return (this->create_ksp(SAME_PRECONDITIONER));
    } 
    inline int    Petsc_::create_ksp(MatStructure flag) {
      PetscErrorCode ierr;

      // - Options
      // SAME_PRECONDITIONER       - Pmat is identical during successive linear solves
      // SAME_NONZERO_PATTERN      - Pmat has same structure
      // DIFFERENT_NONZERO_PATTERN - different structure

      ierr = KSPCreate(PETSC_COMM_WORLD, &(this->ksp));CHKERRQ(ierr);
      ierr = KSPSetOperators(this->ksp, this->A, this->A, flag);CHKERRQ(ierr);

      if (this->is_spd()) {
        ierr = KSPSetType(this->ksp, KSPCG);CHKERRQ(ierr);
      } else {
        ierr = KSPSetType(this->ksp, KSPGMRES);CHKERRQ(ierr);
      }

      ierr = KSPSetTolerances(this->ksp,
                              this->get_tolerence(),
                              PETSC_DEFAULT,
                              PETSC_DEFAULT,
                              PETSC_DEFAULT);CHKERRQ(ierr);
      ierr = KSPSetInitialGuessNonzero(this->ksp, PETSC_TRUE);CHKERRQ(ierr);
      
      return (0);
    }

#undef  __FUNCT__
#define __FUNCT__ "uhm::plugin::Petsc_::delete_ksp"
    inline int    Petsc_::delete_ksp() {
      PetscErrorCode ierr;
      ierr = KSPDestroy(&this->ksp);CHKERRQ(ierr);
      return (0);
    }

#undef  __FUNCT__
#define __FUNCT__ "uhm::plugin::Petsc_::set_pc_on_ksp"
    inline int    Petsc_::set_pc_on_ksp(PCType type) {

      // PCType
      // PCNONE, PCJACOBI, PCSOR, PCLU, PCSHELL, PCBJACOBI, PCMG, PCILU
      // PCASM, PCKSP, PCCOMPOSITE, PCSPI, PCCHOLESKY, PCSAMG, PCHYPRE and so on
      //
      // see petscpc.h
      // Note that PCBJACOBI call PCILU for the block matrix

      PetscErrorCode ierr;
      PC pc;
      ierr = KSPGetPC(this->ksp, &pc);CHKERRQ(ierr);
      ierr = PCSetType(pc, type);CHKERRQ(ierr);
      
      return (0);
    }

#undef  __FUNCT__
#define __FUNCT__ "uhm::plugin::Petsc_::solve_dist"
    inline int    Petsc_::solve_dist() {
      PetscErrorCode ierr;

      // loop through n_rhs
      for (int i=0;i<this->get_n_rhs();++i) {
        
        ierr = this->set_values_on_rhs_begin(this->B, this->b, i);CHKERRQ(ierr);
        ierr = KSPSetFromOptions(this->ksp);CHKERRQ(ierr);
        ierr = KSPSetUp(this->ksp);CHKERRQ(ierr);
        ierr = KSPSolve(this->ksp, this->b, this->x);CHKERRQ(ierr);
        ierr = this->set_values_on_rhs_end(this->X, this->x, i);CHKERRQ(ierr);

        if (this->verbose_col_op) {
          ierr = this->convergence();CHKERRQ(ierr);
        }
      }

      return (0);
    }

#undef  __FUNCT__
#define __FUNCT__ "uhm::plugin::Petsc_::check_dist"
    inline int   Petsc_::check_dist() { 
 
      PetscErrorCode ierr;
      PetscReal norm;

      int n = this->get_n_rhs();
      for (int i=0;i<n;++i) {

        ierr = this->set_values_on_rhs_begin(this->B, this->b, i);CHKERRQ(ierr);
        ierr = this->set_values_on_rhs_begin(this->X, this->x, i);CHKERRQ(ierr);

        ierr = MatMult(this->A, this->x, this->r);CHKERRQ(ierr);
        ierr = VecAXPY(this->r, -1.0, this->b);CHKERRQ(ierr);
        ierr = VecNorm(this->r, NORM_1, &norm);CHKERRQ(ierr);

        ierr = this->set_values_on_rhs_end(this->R, this->r, i);CHKERRQ(ierr);

        if (this->verbose_col_op) {
          ierr = PetscPrintf(PETSC_COMM_WORLD,
                             ">> Norm of residual for column %d is % E\n", 
                             i, norm);CHKERRQ(ierr);
        }
      }
      ierr = MatNorm(this->R, NORM_1, &norm);CHKERRQ(ierr);
      ierr = PetscPrintf(PETSC_COMM_WORLD,
                         ">> Norm of residual is % E\n", 
                         norm);CHKERRQ(ierr);

      return (0);
    }

#undef  __FUNCT__
#define __FUNCT__ "uhm::plugin::Petsc_::condition"
    inline int   Petsc_::condition(PetscReal &cond) {
      return (0);
    }

#undef  __FUNCT__
#define __FUNCT__ "uhm::plugin::Petsc_::view"
    inline int   Petsc_::view(int obj, PetscViewer viewer) {
      PetscErrorCode ierr;

      // PETSC_VIEWER_STDOUT_SELF  - standard output (default)
      // PETSC_VIEWER_STDOUT_WORLD - synchronized standard output 
      // PETSC_VIEWER_DRAW_WORLD   - graphical display of nonzero structure

      switch (obj) {
      case UHM_PLUGIN_PETSC_KSP:   ierr = KSPView(this->ksp, viewer);      break;
      case UHM_PLUGIN_PETSC_MAT_A: ierr = MatView(this->A, viewer);        break;
      case UHM_PLUGIN_PETSC_MAT_B: ierr = MatView(this->B, viewer);        break;
      case UHM_PLUGIN_PETSC_MAT_X: ierr = MatView(this->X, viewer);        break;
      case UHM_PLUGIN_PETSC_MAT_R: ierr = MatView(this->R, viewer);        break;
      case UHM_PLUGIN_PETSC_VEC_B: ierr = VecView(this->b, viewer);        break;
      case UHM_PLUGIN_PETSC_VEC_X: ierr = VecView(this->x, viewer);        break;
      case UHM_PLUGIN_PETSC_VEC_R: ierr = VecView(this->r, viewer);        break;
      default:
        SETERRQ(PETSC_COMM_SELF, 1,
                ">> Argument obj is wrong");
      }
      CHKERRQ(ierr);

      return (0);
    }

#undef  __FUNCT__
#define __FUNCT__ "uhm::plugin::Petsc_::get_mat"
    inline int   Petsc_::get_mat(int obj, Mat *M) {
      PetscErrorCode ierr;

      switch (obj) {
      case UHM_PLUGIN_PETSC_MAT_A: *M = this->B; break;
      case UHM_PLUGIN_PETSC_MAT_B: *M = this->B; break;
      case UHM_PLUGIN_PETSC_MAT_X: *M = this->X; break;
      case UHM_PLUGIN_PETSC_MAT_R: *M = this->R; break;
      default:
        SETERRQ(PETSC_COMM_SELF, 1,
                ">> Argument obj is wrong");
      }
      return (0);
    }

#undef  __FUNCT__
#define __FUNCT__ "uhm::plugin::Petsc_::get_vec"
    inline int   Petsc_::get_vec(int obj, Vec *v) {
      PetscErrorCode ierr;

      switch (obj) {
      case UHM_PLUGIN_PETSC_VEC_B: *v = this->b; break;
      case UHM_PLUGIN_PETSC_VEC_X: *v = this->x; break;
      case UHM_PLUGIN_PETSC_VEC_R: *v = this->r; break;
      default:
        SETERRQ(PETSC_COMM_SELF, 1,
                ">> Argument obj is wrong");
      }
      return (0);
    }

#undef  __FUNCT__
#define __FUNCT__ "uhm::plugin::Petsc_::convergence"
    inline int   Petsc_::convergence() {
      return (this->convergence(PETSC_COMM_WORLD, stdout));
    }
    inline int   Petsc_::convergence(MPI_Comm comm, FILE *stream) {
      PetscErrorCode ierr;
      KSPConvergedReason reason;

      ierr = KSPGetConvergedReason(this->ksp, &reason);CHKERRQ(ierr);
      if (reason < 0) {
        ierr = PetscFPrintf(comm, stream, 
                            ">> DIVERGED,  reason = %d\n\n", 
                            reason);CHKERRQ(ierr);
      } else {
        PetscInt its;
        ierr = KSPGetIterationNumber(this->ksp, &its);CHKERRQ(ierr);
        ierr = PetscFPrintf(comm, stream, 
                            ">> CONVERGED, reason = %d, iteration = %d\n\n", 
                            reason, its);CHKERRQ(ierr);
      }
      
      return (0);
    }

#undef  __FUNCT__
#define __FUNCT__ "uhm::plugin::Petsc_::create_matrix"
    inline int   Petsc_::create_matrix() {

      // - Options
      // -mat_type seqaij   - AIJ type, uses MatCreateSeqAIJ()
      // -mat_type mpiaij   - AIJ type, uses MatCreateMPIAIJ()
      // -mat_type seqdense - dense type, uses MatCreateSeqDense()
      // -mat_type mpidense - dense type, uses MatCreateMPIDense()
      // -mat_type seqbaij  - block AIJ type, uses MatCreateSeqBAIJ()
      // -mat_type mpibaij  - block AIJ type, uses MatCreateMPIBAIJ()

      // - MatType
      // - MATSAME
      // - MATSEQMAIJ, MATMPIMAIJ
      // - MATMAIJ      
      // - MATIS        
      // - MATSEQAIJ, MATMPIAIJ    
      // - MATAIJ **      
      // - MATSHELL     
      // - MATSEQDENSE, MATMPIDENSE  
      // - MATDENSE **    
      // - MATSEQBAIJ, MATMPIBAIJ   

      PetscErrorCode ierr;

      // ** Create A and B matrix
      {
        PetscInt n_rows = this->get_n_dof();

        ierr = MatCreate(PETSC_COMM_WORLD, &(this->A));CHKERRQ(ierr);
        ierr = MatSetSizes(this->A, PETSC_DECIDE,PETSC_DECIDE,
                           n_rows, n_rows);CHKERRQ(ierr);
        ierr = MatSetType(this->A, MATMPIAIJ);CHKERRQ(ierr);
        ierr = MatSetFromOptions(this->A);CHKERRQ(ierr);
        ierr = MatSetUp(this->A);CHKERRQ(ierr);
        
        PetscInt n_cols = this->get_n_rhs();

        ierr = MatCreate(PETSC_COMM_WORLD, &(this->B));CHKERRQ(ierr);
        ierr = MatSetSizes(this->B, PETSC_DECIDE,PETSC_DECIDE,
                           n_rows, n_cols );CHKERRQ(ierr);
        ierr = MatSetType(this->B, MATMPIDENSE);CHKERRQ(ierr);
        ierr = MatSetFromOptions(this->B);CHKERRQ(ierr);
        ierr = MatSetUp(this->B);CHKERRQ(ierr);
      }

      PetscInt own_begin, own_end;
      ierr = MatGetOwnershipRange(this->A, &own_begin, &own_end);CHKERRQ(ierr);

      // ** Assembling schur complements from children
      std::vector< std::pair< Element, Element > > elts;
      std::vector< int > idx;
      std::vector< bool > visit;


      // find index set
      for (int i=0;i<this->get_n_children();++i) {
        Mesh c          = this->get_child(i);
        Element e_local = c->get_root();
        Element e_dist  = this->get_self()->find_element( c->get_id() );
        elts.push_back( std::make_pair( e_local, e_dist ) );
        e_dist->export_diag_matrix(UHM_ABR, idx);
      }

      // Assembly for the preallocation
      if (this->get_n_children() > 1) {
        std::sort  (idx.begin(), idx.end());
        std::unique(idx.begin(), idx.end());
        
        visit.assign(idx.size(), false);
        
        for (int i=0;i<this->get_n_children();++i) {
          
        }        
      } else {

      }

      for (int i=0;i<this->get_n_children();++i) {
        Mesh c          = this->get_child(i);
        Element e_local = c->get_root();
        Matrix hm       = e_local->get_matrix();
        Element e_dist  = this->get_self()->find_element( c->get_id() );

        PetscScalar *buf_A, *buf_X, *buf_T; 
        PetscInt n_rows, n_cols;
        std::vector< PetscInt > idx_rows, idx_cols;

        size_t size;
        ierr = PetscDataTypeGetSize((PetscDataType)this->get_petsc_datatype(),
                                    &size);CHKERRQ(ierr);

        e_dist->export_abr_matrix(idx_rows);
        n_rows = n_cols = idx_rows.size();
        buf_T = (PetscScalar*)hm->get_buffer( UHM_ABR ).first;

        if (this->verbose) {
          ierr = PetscPrintf(PETSC_COMM_WORLD,
                             ">> Begin A setvalue for child %d\n", i);CHKERRQ(ierr);
        }

        ierr = PetscMalloc(n_rows*n_cols*size, &buf_A);CHKERRQ(ierr);
        for (int j=0;j<n_rows;++j) 
          for (int i=0;i<n_cols;++i) 
            buf_A[i*n_rows+j] = buf_T[i+j*n_cols];

        ierr = MatSetValues(this->A,
                            n_rows, &(idx_rows[0]),
                            n_cols, &(idx_rows[0]),
                            buf_A, ADD_VALUES);CHKERRQ(ierr);

        ierr = PetscFree(buf_A);CHKERRQ(ierr);

        if (this->verbose) {
          ierr = PetscPrintf(PETSC_COMM_WORLD,
                             ">> End   A setvalue for child %d\n", i);CHKERRQ(ierr);
        }

        n_cols = this->get_n_rhs();
        idx_cols.reserve( n_cols );
        for (int j=0;j<n_cols;++j)
          idx_cols.push_back(j);
        buf_T = (PetscScalar*)hm->get_buffer( UHM_XB ).first;



        if (this->verbose) {
          ierr = PetscPrintf(PETSC_COMM_WORLD,
                             ">> Begin B setvalue for child %d\n", i);CHKERRQ(ierr);
        }

        ierr = PetscMalloc(n_rows*n_cols*size, &buf_X);CHKERRQ(ierr);
        for (int j=0;j<n_rows;++j)
          for (int i=0;i<n_cols;++i)
            buf_X[i*n_rows+j] = buf_T[i+j*n_cols];

        ierr = MatMPIDenseSetPreallocation(this->B, PETSC_NULL);CHKERRQ(ierr);
        ierr = MatSetValues(this->B,
                            n_rows, &(idx_rows[0]),
                            n_cols, &(idx_cols[0]),
                            buf_X, ADD_VALUES);CHKERRQ(ierr);

        ierr = PetscFree(buf_X);CHKERRQ(ierr);

        if (this->verbose) {
          ierr = PetscPrintf(PETSC_COMM_WORLD,
                             ">> End   B setvalue for child %d\n", i);CHKERRQ(ierr);
        }
      }

      ierr = MatSetOption(this->A,
                          MAT_STRUCTURALLY_SYMMETRIC, PETSC_TRUE);CHKERRQ(ierr);

      // This should be better then begin end and begin end.
      // Ask Victor
      ierr = MatAssemblyBegin(this->A, MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
      ierr = MatAssemblyBegin(this->B, MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
      ierr = MatAssemblyEnd  (this->B, MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
      ierr = MatAssemblyEnd  (this->A, MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);


      //       PetscInt  dd;
      //       PetscBool missing;
      //       ierr = MatMissingDiagonal(this->A, &missing, &dd);CHKERRQ(ierr);
      //       if (missing)
      //         printf(" A miss diagonal at %d\n", dd);

      ierr = MatDuplicate(this->B, MAT_DO_NOT_COPY_VALUES, &(this->X));CHKERRQ(ierr);
      ierr = MatDuplicate(this->X, MAT_DO_NOT_COPY_VALUES, &(this->R));CHKERRQ(ierr);

      // ** Create vector workspace
      ierr = VecCreate(PETSC_COMM_WORLD, &(this->b));CHKERRQ(ierr);
      ierr = VecSetSizes(this->b, PETSC_DECIDE, this->get_n_dof());CHKERRQ(ierr);
      ierr = VecSetFromOptions(this->b);CHKERRQ(ierr);

      ierr = VecDuplicate(this->b, &this->x);CHKERRQ(ierr);
      ierr = VecDuplicate(this->x, &this->r);CHKERRQ(ierr);
      
      return (0);
    }

#undef  __FUNCT__
#define __FUNCT__ "uhm::plugin::Petsc_::delete_matrix"
    inline int   Petsc_::delete_matrix() {
      PetscErrorCode ierr;

      // ** Distribute solutions
      for (int i=0;i<this->get_n_children();++i) {
        Mesh c          = this->get_child(i);
        Element e_local = c->get_root();
        Matrix hm       = e_local->get_matrix();
        Element e_dist  = this->get_self()->find_element( c->get_id() );

        PetscScalar *buf_X;
        PetscInt n_rows, n_cols;
        std::vector< PetscInt > idx_rows, idx_cols;

        e_dist->export_abr_matrix(idx_rows);
        n_rows = idx_rows.size();

        n_cols = this->get_n_rhs();
        idx_cols.reserve( n_cols );
        for (int i=0;i<n_cols;++i)
          idx_cols.push_back(i);

        buf_X = (PetscScalar*)hm->get_buffer( UHM_XB ).first;

        ierr = MatGetValues(this->X,
                            n_rows, &(idx_rows[0]),
                            n_cols, &(idx_cols[0]),
                            buf_X);CHKERRQ(ierr);
      }

      ierr = MatDestroy(&(this->A));CHKERRQ(ierr);
      ierr = MatDestroy(&(this->B));CHKERRQ(ierr);
      ierr = MatDestroy(&(this->R));CHKERRQ(ierr);

      ierr = VecDestroy(&(this->x));CHKERRQ(ierr);
      ierr = VecDestroy(&(this->b));CHKERRQ(ierr);
      ierr = VecDestroy(&(this->r));CHKERRQ(ierr);
      
      return (0);
    }

#undef  __FUNCT__
#define __FUNCT__ "uhm::plugin::Petsc_::set_values_on_rhs_begin"
    inline int   Petsc_::set_values_on_rhs_begin(Mat B, Vec b, int loc) {
      PetscErrorCode ierr;
      PetscInt       
        B_begin_row, B_end_row,
        B_begin_col, B_end_col,
        b_begin,     b_end, 
        size;
      PetscScalar    *buf_B, *buf_b;

      ierr = MatGetOwnershipRange      (B, &B_begin_row, &B_end_row);CHKERRQ(ierr);
      ierr = MatGetOwnershipRangeColumn(B, &B_begin_col, &B_end_col);CHKERRQ(ierr);
      ierr = VecGetOwnershipRange      (b, &b_begin,     &b_end);CHKERRQ(ierr);

      if (B_begin_row != b_begin || B_end_row != b_end) {
        SETERRQ(PETSC_COMM_SELF, 1, 
                ">> Distribution of matrix and vector does not match");
      }
      if (loc < B_begin_col || loc >= B_end_col) {
        printf("begin %d end %d loc %d\n", B_begin_col, B_end_col, loc);
        SETERRQ(PETSC_COMM_SELF, 1, 
                ">> Given column index is out of range");
      }

      // below is equivalent statement from petsc but it is slow and 
      // we already know that B is dense
      // ( MatGetColumnVector do search for sparse matrix )
      // ierr = MatGetColumnVector(B, b, loc);CHKERRQ(ierr);

      ierr = MatGetArray(B, &buf_B);CHKERRQ(ierr);
      ierr = VecGetArray(b, &buf_b);CHKERRQ(ierr);

      size = b_end - b_begin;
      for (int i=0;i<size;++i) 
        buf_b[i] = buf_B[loc*size+i];
      
      ierr = VecRestoreArray(b, &buf_b);CHKERRQ(ierr);
      ierr = MatRestoreArray(B, &buf_B);CHKERRQ(ierr);

      // ** Do I need to reset the initial guess
      // zero ? or b ? or random ?
      //ierr = VecCopy(b, x);CHKERRQ(ierr);

      return (0);
    }

#undef  __FUNCT__
#define __FUNCT__ "uhm::plugin::Petsc_::set_values_on_rhs_end"
    inline int   Petsc_::set_values_on_rhs_end(Mat B, Vec b, int loc) {
      PetscErrorCode ierr;
      PetscInt       
        B_begin_row, B_end_row, 
        B_begin_col, B_end_col, 
        b_begin,     b_end, 
        size;
      PetscScalar    *buf_B, *buf_b;


      ierr = MatGetOwnershipRange      (B, &B_begin_row, &B_end_row);CHKERRQ(ierr);
      ierr = MatGetOwnershipRangeColumn(B, &B_begin_col, &B_end_col);CHKERRQ(ierr);
      ierr = VecGetOwnershipRange      (b, &b_begin,     &b_end);CHKERRQ(ierr);

      if (B_begin_row != b_begin || B_end_row != b_end) {
        SETERRQ(PETSC_COMM_SELF, 1,
                ">> Distribution of matrib and vector does not match");
      }
      if (loc < B_begin_col || loc >= B_end_col) {
        SETERRQ(PETSC_COMM_SELF, 1,
                ">> Given column index is out of range");
      }

      ierr = MatGetArray(B, &buf_B);CHKERRQ(ierr);
      ierr = VecGetArray(b, &buf_b);CHKERRQ(ierr);

      size = b_end - b_begin;
      for (int i=0;i<size;++i)
        buf_B[loc*size+i] = buf_b[i];

      ierr = VecRestoreArray(b, &buf_b);CHKERRQ(ierr);
      ierr = MatRestoreArray(B, &buf_B);CHKERRQ(ierr);

      return (0);
    }
    
    inline void Petsc_::lock() {
      UHM_ERROR(!this->is_locked(),">> uhm::plugin::Petsc_ is already locked");

      Mesh m = this->get_self();
      if (!m->is_locked()) m->lock();

      Element root = m->get_root();
      this->n_dof = root->get_n_dof().first;
      root->set_nodes_kind(UHM_NODE_KIND_BOUNDARY);


      int 
        size = this->get_comm_size(), 
        me = this->get_comm_rank(),
        n = root->get_n_children();
      
      this->reset_children( n/size + 1 );

      for (int i=0;i<n;++i) 
        if (i%size == me) 
          this->children.push_back( Mesh_( root->get_child(i)->get_id() ) );

      for (int i=0;i<this->get_n_children();++i) {
        Mesh c = this->get_child(i);
        
        c->insert_element( m->find_element(c->get_id()) );
        c->adjust_element_numbering();
      }

      this->locker = true;

      m->remove_leaves();
      m->remove_non_boundary_nodes_from_leaves();

      root->set_nodes_kind(UHM_NODE_KIND_DEFAULT);
      m->unlock();
      m->lock();

    }
    inline void Petsc_::unlock() { this->locker = false; }
  }
}


#endif
